import os
import cv2
import time
import base64
from flask import Flask
from flask import request
from flask import render_template

import utils
import log_handler


app = Flask(__name__)
model_path = "model/realesrgan_jit.pt"
model = utils.RealESRGANer(4, model_path)

logger = log_handler.get_logger('log.txt')


def resize_img_keep_ratio(img, max_size=512):
    old_size = img.shape[0:2]
    ratio = min(float(max_size) / (old_size[i]) for i in range(len(old_size)))
    ratio = min(1, ratio)
    new_size = tuple([int(i * ratio) for i in old_size])
    img = cv2.resize(img, (new_size[1], new_size[0]))
    logger.info("input old_size: " + str(old_size))
    logger.info("input new_size: " + str(new_size))
    return img


@app.route('/sr', methods=['POST'])
def sr():
    # 接收图片
    upload_file = request.files['file']
    # 获取图片名
    file_name = upload_file.filename
    # 文件保存目录（桌面）
    logger.info(request)
    logger.info("file_name: " + file_name)
    file_path = 'image'
    if upload_file:
        # 地址拼接
        file_path = os.path.join(file_path, file_name)
        # 保存接收的图片到桌面
        # img = base64.b64decode(upload_file)
        upload_file.save(file_path)
        img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

        start_time = time.time()
        img = resize_img_keep_ratio(img)
        img_sr, img_mode = model.enhance(img)
        logger.info("output size: " + str(img_sr.shape[0:2]))
        logger.info("cost time: " + str(time.time() - start_time))

        file, extension = os.path.splitext(file_path)
        output_path = file + "_sr" + extension
        cv2.imwrite(output_path, img_sr)

        with open(output_path, 'rb') as f:
            res = base64.b64encode(f.read())
            return res


@app.route('/')
def index():
    return render_template('hello.html')


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8082)
